//
// Created by inter on 2019-11-07.
//

#include <cstring>
#include <iostream>
#include <sys/stat.h>
#include "aes.h"
#include "aes_impl.h"

int AES::calculateOutputLength(unsigned int input_size, bool is_encrypt) {
    if (is_encrypt) {
        return (input_size / 16 + 1) * 16;
    } else {
        return input_size / 2;
    }
}

bool AES::cbc256Encrypt(unsigned char *input_buf, unsigned int input_size, unsigned char *iv, unsigned char *key,
                        unsigned char *output_buf, unsigned int output_size) {

    if (output_size < (input_size / 16 + 1) * 16)
        return false;

    memset(output_buf, 0, (size_t) output_size);
    memcpy(output_buf, input_buf, (size_t) input_size);

    AES_ctx ctx;
    AES_init_ctx_iv(&ctx, key, iv);
    AES_CBC_encrypt_buffer(&ctx, output_buf, (input_size / 16 + 1) * 16);

    return true;
}

bool AES::cbc256Decrypt(unsigned char *input_buf, unsigned int input_size, unsigned char *iv, unsigned char *key,
                        unsigned char *output_buf, unsigned int output_size) {

    memset(output_buf, 0, (size_t) output_size);
    memcpy(output_buf, input_buf, (size_t) input_size);

    AES_ctx ctx;
    AES_init_ctx_iv(&ctx, key, iv);
    AES_CBC_decrypt_buffer(&ctx, output_buf, output_size);

    return true;
}

bool AES::cbc256EncryptFile(string input_file, string output_file, unsigned char *iv, unsigned char *key) {

    unsigned long file_size = 0;
    FILE *file_fd = fopen(input_file.c_str(), "rb");
    if (file_fd == nullptr) {
        cout << "ERROR: can't open file {"<< input_file <<"}\n" << endl;
        return false;
    } else {
        fseek(file_fd, 0, SEEK_END);
        file_size = (unsigned long)ftell(file_fd);
        fseek(file_fd, 0, SEEK_SET);
    }

    FILE *out_fd = fopen(output_file.c_str(), "wb");
    if (out_fd == nullptr) {
        cout << "ERROR: can't create file {"<< output_file <<"}\n" << endl;
        fclose(file_fd);
        return false;
    }

    unsigned long out_len = (file_size / 16 + 1) * 16;
    auto* file_buf = new unsigned char[out_len]{0};


    AES_ctx ctx;
    AES_init_ctx_iv(&ctx, key, iv);
    unsigned char input_buf[64];
    int read_len , buf_pos = 0;

    while (true) {
        read_len = fread(input_buf, 1, 64, file_fd);
        if (read_len > 0) {
            memcpy(file_buf + buf_pos, input_buf, read_len);
            buf_pos += read_len;
        } else {
            break;
        }
    }

    AES_CBC_encrypt_buffer(&ctx, (unsigned char *) file_buf, out_len);
    fwrite(file_buf, 1, out_len, out_fd);
    fflush(out_fd);

    delete[](file_buf);
    fclose(out_fd);
    fclose(file_fd);
    return true;
}

bool AES::cbc256DecryptFile(string input_file, string output_file, unsigned char* iv, unsigned char* key) {
    unsigned long file_size = 0;
    FILE *file_fd = fopen(input_file.c_str(), "rb");
    if (file_fd == nullptr) {
        return false;
    } else {
        fseek(file_fd, 0, SEEK_END);
        file_size = (unsigned long)ftell(file_fd);
        fseek(file_fd, 0, SEEK_SET);
    }

    FILE *out_fd = fopen(output_file.c_str(), "wb");
    if (out_fd == nullptr) {
        fclose(file_fd);
        return false;
    }

    unsigned long out_len = file_size;
    auto* file_buf = new unsigned char[out_len]{0};


    AES_ctx ctx;
    AES_init_ctx_iv(&ctx, key, iv);
    unsigned char input_buf[64];
    int read_len , buf_pos = 0;

    while (true) {
        read_len = fread(input_buf, 1, 64, file_fd);
        if (read_len > 0) {
            memcpy(file_buf + buf_pos, input_buf, read_len);
            buf_pos += read_len;
        } else {
            break;
        }
    }

    AES_CBC_decrypt_buffer(&ctx, (unsigned char *) file_buf, out_len);
    fwrite(file_buf, 1, (out_len > 16 && file_buf[out_len - 16] == 0x00)? out_len - 16 : out_len, out_fd);
    fflush(out_fd);

    delete[](file_buf);
    fclose(out_fd);
    fclose(file_fd);
    return true;
}